{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat Context Dependency Injection\n",
    "\n",
    "In this tutorial, we’ll build upon the concepts introduced in the [Tools with Dependency Injection](https://github.com/ag2ai/ag2/blob/main/notebook/tools_dependency_injection.ipynb) notebook to demonstrate how to use `ChatContext` for more advanced workflows.\n",
    "\n",
    "By leveraging `ChatContext`, we can track the flow of conversations, ensuring proper function execution order. For example, before retrieving a user’s account balance, we’ll ensure the user has logged in first. This approach prevents unauthorized actions and enhances security.\n",
    "\n",
    "**Benefits of Using `ChatContext`**\n",
    "- Flow Control: It helps enforce the correct sequence of function calls.\n",
    "- Enhanced Security: Ensures actions depend on preconditions like authentication.\n",
    "- Simplified Debugging: Logs conversation history, making it easier to trace issues.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "To install `AG2`, simply run the following command:\n",
    "\n",
    "```bash\n",
    "pip install -U ag2[openai]\n",
    "```\n",
    "\n",
    "> **Note:** If you have been using `autogen` or `ag2`, all you need to do is upgrade it using:  \n",
    "> ```bash\n",
    "> pip install -U autogen\n",
    "> ```\n",
    "> or  \n",
    "> ```bash\n",
    "> pip install -U ag2\n",
    "> ```\n",
    "> as `autogen`, and `ag2` are aliases for the same PyPI package.  \n",
    "\n",
    "\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Annotated, Literal\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from autogen.agentchat import AssistantAgent, UserProxyAgent\n",
    "from autogen.tools.dependency_injection import BaseContext, ChatContext, Depends"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define BaseContext Class\n",
    "\n",
    "The following `BaseContext` class and helper functions are adapted from the previous tutorial. They define the structure for securely handling account data and operations like login and balance retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Account(BaseContext, BaseModel):\n",
    "    username: str\n",
    "    password: str\n",
    "    currency: Literal[\"USD\", \"EUR\"] = \"USD\"\n",
    "\n",
    "\n",
    "alice_account = Account(username=\"alice\", password=\"password123\")\n",
    "bob_account = Account(username=\"bob\", password=\"password456\")\n",
    "\n",
    "account_ballace_dict = {\n",
    "    (alice_account.username, alice_account.password): 300,\n",
    "    (bob_account.username, bob_account.password): 200,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper Functions\n",
    "\n",
    "These functions validate account credentials and retrieve account balances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _verify_account(account: Account):\n",
    "    if (account.username, account.password) not in account_ballace_dict:\n",
    "        raise ValueError(\"Invalid username or password\")\n",
    "\n",
    "\n",
    "def _get_balance(account: Account):\n",
    "    _verify_account(account)\n",
    "    return f\"Your balance is {account_ballace_dict[(account.username, account.password)]}{account.currency}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent Configuration\n",
    "\n",
    "Configure the agents for the interaction.\n",
    "\n",
    "- `config_list` defines the LLM configurations, including the model and API key.\n",
    "- `UserProxyAgent` simulates user inputs without requiring actual human interaction (set to `NEVER`).\n",
    "- `AssistantAgent` represents the AI agent, configured with the LLM settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = [{\"model\": \"gpt-5-nano\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}]\n",
    "agent = AssistantAgent(\n",
    "    name=\"agent\",\n",
    "    llm_config={\"config_list\": config_list},\n",
    ")\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy_1\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    llm_config=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Injecting a ChatContext Parameter\n",
    "\n",
    "Now let’s  upgrade the example from the previous tutorial by introducing the `ChatContext` parameter. This enhancement allows us to enforce proper execution order in the workflow, ensuring that users log in before accessing sensitive data like account balances.\n",
    "\n",
    "The following functions will be registered:\n",
    "\n",
    "- `login`: Verifies the user’s credentials and ensures they are logged in.\n",
    "- `get_balance`: Retrieves the account balance but only if the user has successfully logged in first.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@user_proxy.register_for_execution()\n",
    "@agent.register_for_llm(description=\"Login\")\n",
    "def login(\n",
    "    account: Annotated[Account, Depends(bob_account)],\n",
    ") -> str:\n",
    "    _verify_account(account)\n",
    "    return \"You are logged in\"\n",
    "\n",
    "\n",
    "@user_proxy.register_for_execution()\n",
    "@agent.register_for_llm(description=\"Get balance\")\n",
    "def get_balance(\n",
    "    account: Annotated[Account, Depends(bob_account)],\n",
    "    chat_context: ChatContext,\n",
    ") -> str:\n",
    "    _verify_account(account)\n",
    "\n",
    "    # Extract the list of messages exchanged with the first agent in the conversation.\n",
    "    # The chat_context.chat_messages is a dictionary where keys are agents (objects)\n",
    "    # and values are lists of message objects. We take the first value (messages of the first agent).\n",
    "    messages_with_first_agent = list(chat_context.chat_messages.values())[0]\n",
    "\n",
    "    login_function_called = False\n",
    "    for message in messages_with_first_agent:\n",
    "        if \"tool_calls\" in message and message[\"tool_calls\"][0][\"function\"][\"name\"] == \"login\":\n",
    "            login_function_called = True\n",
    "            break\n",
    "\n",
    "    if not login_function_called:\n",
    "        raise ValueError(\"Please login first\")\n",
    "\n",
    "    balance = _get_balance(account)\n",
    "    return balance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we initiate a chat to retrieve the balance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(agent, message=\"Get users balance\", max_turns=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Chat Context Dependency Injection",
   "tags": [
    "tools",
    "dependency injection",
    "function calling"
   ]
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
